"""
RAG 系统独立入口
支持文档存储和查询功能
"""

import sys
import os
import json
import argparse
from typing import List, Dict, Any

# 添加项目根目录到 Python 路径
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, project_root)

# 切换到项目根目录
os.chdir(project_root)

from base.config import config
from base.logger import logger, setup_logger
from core.rag_system import RAGSystem


class RAGMain:
    """RAG 系统主入口"""
    
    def __init__(self):
        # 初始化日志
        setup_logger()
        logger.info("RAG 系统启动")
        
        # 初始化 RAG 系统
        self.rag_system = RAGSystem()
    
    def add_documents_from_file(self, file_path: str) -> bool:
        """
        从文件添加文档
        
        Args:
            file_path: 文档文件路径（JSON 格式）
            
        Returns:
            是否添加成功
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                documents = json.load(f)
            
            if not isinstance(documents, list):
                logger.error("文档文件格式错误，应为文档列表")
                return False
            
            return self.rag_system.add_documents(documents)
            
        except Exception as e:
            logger.error(f"从文件添加文档失败: {e}")
            return False
    
    def add_documents_interactive(self):
        """交互式添加文档"""
        print("=== 交互式添加文档 ===")
        documents = []
        
        while True:
            print("\n请输入文档信息（输入 'done' 完成添加）:")
            
            doc_id = input("文档ID: ").strip()
            if doc_id.lower() == 'done':
                break
            
            content = input("文档内容: ").strip()
            if not content:
                print("文档内容不能为空")
                continue
            
            # 生成示例嵌入向量（实际应用中需要使用真实的嵌入模型）
            import random
            dimension = config.getint('embedding', 'dimension', fallback=768)
            embedding = [random.random() for _ in range(dimension)]
            
            metadata = {}
            meta_input = input("元数据（JSON格式，可选）: ").strip()
            if meta_input:
                try:
                    metadata = json.loads(meta_input)
                except json.JSONDecodeError:
                    print("元数据格式错误，将使用空元数据")
                    metadata = {}
            
            document = {
                'doc_id': doc_id,
                'content': content,
                'embedding': embedding,
                'metadata': metadata
            }
            
            documents.append(document)
            print(f"已添加文档: {doc_id}")
        
        if documents:
            success = self.rag_system.add_documents(documents)
            if success:
                print(f"\n成功添加 {len(documents)} 个文档到知识库")
            else:
                print("\n文档添加失败")
        else:
            print("\n没有添加任何文档")
    
    def query_interactive(self):
        """交互式查询"""
        print("=== RAG 查询系统 ===")
        print("输入 'quit' 退出，'clear' 清空历史，'stats' 查看统计信息")
        
        while True:
            question = input("\n请输入您的问题: ").strip()
            
            if question.lower() == 'quit':
                break
            elif question.lower() == 'clear':
                self.rag_system.clear_history()
                print("对话历史已清空")
                continue
            elif question.lower() == 'stats':
                stats = self.rag_system.get_system_stats()
                print(f"系统统计信息: {json.dumps(stats, indent=2, ensure_ascii=False)}")
                continue
            elif not question:
                print("问题不能为空")
                continue
            
            # 执行查询
            result = self.query_single(question)
            
            # 显示结果
            if result['success']:
                print(f"\n答案: {result['answer']}")
                print(f"\n相关文档数量: {len(result['similar_docs'])}")
                print(f"查询类别: {result['query_features']['category']}")
            else:
                print(f"\n查询失败: {result['answer']}")
    
    def query_single(self, question: str, use_history: bool = True) -> Dict[str, Any]:
        """
        单次查询
        
        Args:
            question: 用户问题
            use_history: 是否使用对话历史
            
        Returns:
            查询结果
        """
        return self.rag_system.query(question, use_history)
    
    def batch_query_from_file(self, file_path: str, output_path: str = None):
        """
        从文件批量查询
        
        Args:
            file_path: 问题文件路径（每行一个问题）
            output_path: 结果输出路径
        """
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                questions = [line.strip() for line in f if line.strip()]
            
            results = []
            for i, question in enumerate(questions, 1):
                print(f"处理问题 {i}/{len(questions)}: {question}")
                result = self.query_single(question, use_history=False)
                results.append(result)
            
            # 保存结果
            if output_path:
                with open(output_path, 'w', encoding='utf-8') as f:
                    json.dump(results, f, indent=2, ensure_ascii=False)
                print(f"结果已保存到: {output_path}")
            else:
                print(json.dumps(results, indent=2, ensure_ascii=False))
                
        except Exception as e:
            logger.error(f"批量查询失败: {e}")
    
    def show_system_info(self):
        """显示系统信息"""
        print("=== RAG 系统信息 ===")
        
        # 配置信息
        print("\n配置信息:")
        print(f"Milvus 主机: {config.get('milvus', 'host', fallback='localhost')}")
        print(f"Milvus 端口: {config.get('milvus', 'port', fallback='19530')}")
        print(f"Redis 主机: {config.get('redis', 'host', fallback='localhost')}")
        print(f"Redis 端口: {config.get('redis', 'port', fallback='6379')}")
        print(f"嵌入维度: {config.get('embedding', 'dimension', fallback='768')}")
        
        # 系统统计
        stats = self.rag_system.get_system_stats()
        print(f"\n系统统计: {json.dumps(stats, indent=2, ensure_ascii=False)}")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='EduRAG 系统')
    parser.add_argument('--mode', choices=['add', 'query', 'batch', 'info'], 
                       default='query', help='运行模式')
    parser.add_argument('--file', help='输入文件路径')
    parser.add_argument('--output', help='输出文件路径')
    
    args = parser.parse_args()
    
    try:
        rag_main = RAGMain()
        
        if args.mode == 'add':
            if args.file:
                success = rag_main.add_documents_from_file(args.file)
                if success:
                    print("文档添加成功")
                else:
                    print("文档添加失败")
            else:
                rag_main.add_documents_interactive()
        
        elif args.mode == 'query':
            rag_main.query_interactive()
        
        elif args.mode == 'batch':
            if not args.file:
                print("批量查询模式需要指定输入文件")
                return
            rag_main.batch_query_from_file(args.file, args.output)
        
        elif args.mode == 'info':
            rag_main.show_system_info()
    
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    except Exception as e:
        logger.error(f"程序运行错误: {e}")
        print(f"程序运行错误: {e}")


if __name__ == "__main__":
    main()